import numpy as np
import torch
import cv2
from PIL import Image
import torchvision.transforms as transforms
from torch.utils.data import Dataset
import os
import random

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
)


class MyData(Dataset):
    def __init__(self, x, y, h=416, w=416, transform=None):
        super(MyData, self).__init__()
        self.x = x
        self.y = y
        self.transform = transform
        self.h = h
        self.w = w

    def __getitem__(self, index):
        img = self.x[index]
        label = self.y[index]
        img = Image.fromarray(img)
        # 改变分辨率进行实验############################
        img = img.resize((self.h, self.w))
        img = self.get_random_data(img, (self.h, self.w))
        #############################################
        img = img / 255.0
        if transform is not None:
            img = transform(img)
        return img, label

    def __len__(self):
        return len(self.x)

    def rand(self, a=0, b=1):
        return np.random.rand() * (b - a) + a

    def get_random_data(self, image, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5, random=True):
        """实时数据增强的随机预处理"""
        iw, ih = image.size
        h, w = input_shape

        if not random:
            scale = min(w / iw, h / ih)
            nw = int(iw * scale)
            nh = int(ih * scale)
            dx = (w - nw) // 2
            dy = (h - nh) // 2

            image = image.resize((nw, nh), Image.BICUBIC)
            new_image = Image.new('RGB', (w, h), (128, 128, 128))
            new_image.paste(image, (dx, dy))
            image_data = np.array(new_image, np.float32)
            return image_data

        # 是否翻转图片
        flip = self.rand() < .5
        if flip:
            image = image.transpose(Image.FLIP_LEFT_RIGHT)

        # 是否旋转图片
        rotate = self.rand() < .5
        if rotate:
            image = image.rotate(90)

        # 色域变换
        hue = self.rand(-hue, hue)
        sat = self.rand(1, sat) if self.rand() < .5 else 1 / self.rand(1, sat)
        val = self.rand(1, val) if self.rand() < .5 else 1 / self.rand(1, val)
        x = cv2.cvtColor(np.array(image, np.float32) / 255, cv2.COLOR_RGB2HSV)
        x[..., 0] += hue * 360
        x[..., 0][x[..., 0] > 1] -= 1
        x[..., 0][x[..., 0] < 0] += 1
        x[..., 1] *= sat
        x[..., 2] *= val
        x[x[:, :, 0] > 360, 0] = 360
        x[:, :, 1:][x[:, :, 1:] > 1] = 1
        x[x < 0] = 0
        image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB) * 255
        return image_data


def fun(x, a, b, c=None):
    if x <= a:
        return 0
    elif x >= a + 1 and x <= b:
        return 1
    if c is None:
        if x >= b + 1:
            return 2
    else:
        if x <= c and x >= b + 1:
            return 2
        elif x > c:
            return 3


import shutil


def GenData(path, h=416, w=416, a=859, b=1715):
    """
    :param path: the directory of raw pictures
    :param batch_size:
    :return: train_load
    """
    num_pics = len(os.listdir(path))
    data = np.zeros((num_pics, h,w, 3), dtype=np.uint8)
    index = list(range(num_pics))
    # label = np.array(list(map(lambda x: fun(x, 2667), index)))
    # label = label.reshape(-1, 1)
    # label = list(map(lambda x: fun(x,80,106), index))
    # label = list(map(lambda x: fun(x, 72,95), index))
    # label = list(range(0, num_pics))
    for i, filename in enumerate(os.listdir(path)):
        if filename.endswith(".jpg"):
            img = cv2.imread(os.path.join(path, filename))
            img = cv2.resize(img, (h, w))
            data[i] = img
    label = list(map(lambda x: fun(x, a, b), index))
    label = torch.Tensor(label)
    # shuffle 数据
    index = list(range(0, num_pics))
    random.shuffle(index)
    data = data[index]
    label = label[index]
    return data, label
